# This script defines function to automize the regressions.

run_reg = function(models, print=T) {
  # This function runs a linear regression with robust standard errors for every
  # model specified in `models` (specified in formula, subset).
  # It returns the model list with the estimated models ($reg) and standard errors ($se).
  
  for (type in names(models)) {
    for (mod_name in names(models[[type]]$forms)) {
      # Print currently estimated model
      if(print) {print(paste("Estimating", type, mod_name, "..."))}
      
      # Load data set and formula
      datf = as.data.frame(get(models[[type]]$data_name)) 
      formu = models[[type]]$forms[[mod_name]]
      
      # Attach weights
      if (is.null(models[[type]]$wgts[[mod_name]])) {
        datf$wgt_reg = rep(1, nrow(datf))
      } else {
        datf$wgt_reg = models[[type]]$wgts[[mod_name]]
      }
        
      # Should only a subset be considered?
      subset = models[[type]]$subset[[mod_name]]
      if (is.null(subset)) {subset = rep(T, nrow(datf))}
      datf = datf[subset, ]
      
      # Estimate model
      mod = lm(formula = formu, data = datf, weights = wgt_reg, na.action = na.omit)
      models[[type]]$reg[[mod_name]] = mod
      
      # Derive number of clusters
      models[[type]]$reg[[mod_name]]$num_cluster = ifelse(
        models[[type]]$teq_se == "cse",
        length(unique(datf[names(mod$fitted.values), models[[type]]$cluster_name])),
        NA
      )
      
      # Derive dgf
      models[[type]]$reg[[mod_name]]$dgf = ifelse(
        models[[type]]$teq_se == "cse",
        models[[type]]$reg[[mod_name]]$num_cluster - (length(mod$coefficients)-1),
        summary(mod)$df[2]
      )
      
      # Estimate covariance matrix
      if (models[[type]]$teq_se == "rse") {
        # Robust standard errors
        cov = vcovHC(mod, type = "HC1")
      } else if (models[[type]]$teq_se == "cse") {
        # Clustered standard errors
        cov = cluster.vcov(mod, cluster = datf[,models[[type]]$cluster_name], df_correction = T)
      }
      models[[type]]$vcov[[mod_name]] = cov
      
      # Derive standard errors (save twice, once in regression model,
      # once in model list for stargazer)
      models[[type]]$se[[mod_name]] = sqrt(diag(cov))
      models[[type]]$reg[[mod_name]]$se = sqrt(diag(cov))
      
      # Derive pvalues
      models[[type]]$reg[[mod_name]]$p = pvalues(models[[type]]$reg[[mod_name]])
    }
  }
  
  return(models)
}


pvalues = function(mod, beta=0) {
  # Derives the p-value (based on a standard normal approximation) for
  # all coefficients in a mod.
  # NOTE: Applies to a regression model, not the full model.
  
  # retrieve coefficients, se, dgf
  coeff = mod$coefficients
  se = mod$se
  dgf = mod$dgf
  
  # derive t statistics
  t = (coeff - beta) / se
  
  # derive p values
  p = 2*(1 - pt(abs(t),dgf))
  
  # save in table
  tst = rbind(coeff,t,p)
  
  # transpose and save as data frame
  tst = data.frame(t(tst))
  colnames(tst) = c("coeff", "t", "p")
  
  return(tst)
}


print_reg = function(models) {
  # This function uses stargazer to print all regressions in models to latex code.
  # Parameters for the stargazer function are specified in models.
  
  # Run stargazer
  for (type in names(models)) {
    if (is.null(models[[type]]$keep)) {models[[type]]$keep = NULL}
    if (is.null(models[[type]]$order)) {models[[type]]$order = NULL}
    if (is.null(models[[type]]$single.row)) {models[[type]]$single.row = F}
    
    # Produce stargazer table
    tab = capture.output(stargazer(
      models[[type]]$reg,
      se = models[[type]]$se,
      p = models[[type]]$p,
      title = sub("_", "", models[[type]]$title),
      style = "aer",
      column.labels = models[[type]]$column.labels,
      covariate.labels = models[[type]]$covariate.labels,
      order = models[[type]]$order,
      dep.var.caption = sub("_", "", models[[type]]$dep.var.caption),
      dep.var.labels = models[[type]]$dep.var.labels,
      label = models[[type]]$label,
      omit = models[[type]]$omit.var,
      omit.labels = models[[type]]$omit.labels,
      keep = models[[type]]$keep,
      add.lines = models[[type]]$add.lines,
      omit.yes.no = c("Yes", "No"),
      df = F,
      omit.stat = c("adj.rsq", "f", "res.dev", "ser"),
      font.size = models[[type]]$font.size,
      notes = models[[type]]$notes,
      digits = models[[type]]$digits,
      float = ifelse(is.null(models[[type]]$float),T,models[[type]]$float),
      omit.table.layout = models[[type]]$omit.table.layout,
      single.row = models[[type]]$single.row
    ))
    cat(tab, sep = '\n', file = paste0(path_to("tables"), type, ".tex"))
  }
  
  # Add tabularx environment
  for (type in names(models)) {
    if (!is.null(models[[type]]$tabularx)) {
      # Load the original table
      tab = paste(readLines(paste0(path_to("tables"), type, ".tex")), collapse="\n")
      ncol = length(models[[type]]$forms)
      
      # Make replacements
      tab = gsub("\\{tabular\\}", "\\{tabularx\\}", tab)
      tab = gsub("@\\{\\\\extracolsep\\{5pt\\}", "\\\\linewidth", tab)
      tab = gsub(
        paste0("\\}l", paste0(rep("c", ncol), collapse=""), "\\}"),
        paste0("\\}\\{X", paste0(rep("Y", ncol), collapse=""), "\\}"),
        tab
      )
      tab = paste0("\\newcolumntype{Y}{>{\\centering\\arraybackslash}X}\n", tab)
      
      # Save table
      write(tab, file = paste0(path_to("tables"), type, ".tex"))
    }
  }
  
  
  # Add joint column labels
  for (type in names(models)) {
    if (!is.null(models[[type]]$custom.joint.col.labels)) {
      # Split the original table
      tab = paste(readLines(paste0(path_to("tables"), type, ".tex")), collapse="\n")
      splitter = paste0("\\\\[-1.8ex] & \\multicolumn{", length(models[[type]]$forms), "}{c}{}", " \\\\ \n")
      splitted = strsplit(tab, split=splitter, fixed=T)[[1]]
      
      # Warning if there are more than two splitted object
      if (length(splitted) != 2) {warning("Step \"Add joint column labels\" fails. Cannot locate line.")}
      
      # Rejoin the table with the new line added
      tab = paste0(
        splitted[1],
        splitter,
        models[[type]]$custom.joint.col.labels,
        splitted[2]
      )
      
      # Save table
      write(tab, file = paste0(path_to("tables"), type, ".tex"))
    }
  }
  
  # Add title for add.lines
  for (type in names(models)) {
    if (!is.null(models[[type]]$add.lines.title)) {
      # How do the additional lines look like?
      add.lines = models[[type]]$add.lines
      last_line = add.lines[[length(add.lines)]]
      
      # Split the original table
      tab = paste(readLines(paste0(path_to("tables"), type, ".tex")), collapse="\n")
      splitter = "& \\\\ "
      splitted = strsplit(tab, split=splitter, fixed=T)[[1]]
      length_splitted = length(splitted)
      
      # Rejoin the table with the new line added
      tab = paste0(
        paste0(splitted[1:length_splitted-1],collapse=splitter),
        splitter,
        paste0("\n\\multicolumn{",length(last_line),"}{l}{", models[[type]]$add.lines.title,"}  \\\\\n"),
        splitted[length_splitted]
      )
      
      # Split the original table again
      splitter = paste0(paste0(last_line, collapse=" & "), " \\\\")
      splitted = strsplit(tab, split=splitter, fixed=T)[[1]]
      
      # Warning if there are more than two splitted object
      if (length(splitted) != 2) {warning("Step \"Add title for add.lines\" fails. Cannot locate line.")}
      
      # Rejoin the table with the new line added
      tab = paste0(
        splitted[1],
        splitter,
        " \n\\\\ \n",
        splitted[2]
      )
      
      # Save table
      write(tab, file = paste0(path_to("tables"), type, ".tex"))
    }
  }
}

panel_table = function(models, p_titles, name, font.size="tiny") {
  # Takes the stargazer output (which needs to exist already) and creates a panel
  # table. The titles of the panels are specified in *p_titles*.
  
  # Define building blocks
  tab_start = paste0(
    "\\newcolumntype{Y}{>{\\centering\\arraybackslash}X}\n",
    "\\centering\n\\", font.size, "\n"
  )
  tabx_end = "\\\\[0ex]\n\\end{tabularx}\n\n"
  tabx_end_last = "\\\\[0ex]\n\\hline\\hline\n\\end{tabularx}\n\n"
  
  # Create table  
  tab = tab_start
  for (type in names(models)) {
    # Import the stand-alone table of the regression.
    new_panel = paste(readLines(paste0(path_to("tables"), type, ".tex")), collapse="\n")
    # substract stargazer's first lines
    new_panel = sub(".*\\\\hline \\n\\\\hline \\\\\\\\\\[-1\\.8ex\\] \\n", "", new_panel)
    # substract stargazer's last lines
    new_panel = sub("\\\\hline \\\\\\\\\\[-1\\.8ex\\] \\n\\\\textit.*", "", new_panel)
    
    # Define a new building block
    tabx_start = paste0(
      "\\begin{tabularx}{\\linewidth}{X",
      paste0(rep("Y", length(models[[type]]$forms)), collapse=""),
      "}\n\\hline\\hline\\\\\n\\multicolumn{",
      length(models[[type]]$forms),
      "}{l}{\\textbf{",
      p_titles[which(type == names(models))],
      "}} \\\\\n\\midrule\n"
    )
    
    # Attach to tab. (special ending for last panel)
    if (type != names(models)[length(names(models))]) {
      tab = paste0(tab, tabx_start, new_panel, tabx_end)
    } else {
      tab = paste0(tab, tabx_start, new_panel, tabx_end_last)
    }
  }
  
  # Save table
  write(tab, path_to("tables", name))
}

reg_latex = function(models, file_name) {
  # Runs the regressions (run_reg), prints the latex code (print_reg),
  # and prints latex code to combine all models (combine_latex).
  # Returns models with the regressions.
  
  models = run_reg(models)
  print_reg(models)
  
  return(models)
}



compile_table = function(file_title, open=F) {
  # Compiles the table in a temporary tex file and opens it.
  
  # Import table_viewer template, insert file path, and write to file.
  tex_file = path_to("table_viewer", file_title)
  tex = readLines(path_to("table_viewer_temp"))
  tex[grepl("INSERT FILE PATH HERE", tex)] = paste0("\\input{../", file_title, "}")
  writeLines(tex, tex_file)
  
  # Latex: compile
  source(path_to("function_latex"))
  compile_latex(file = tex_file, folder="table_viewer")
  
  # Open compiled pdf
  if (open) {system(paste0('open "', sub(".tex", ".pdf", tex_file), '"'))}
}
